import rclpy
from rclpy.node import Node
from sensor_msgs.msg import JointState
from std_msgs.msg import Int8, Float32
from std_srvs.srv import SetBool
from robot_drivers.msg import Jointpos
from robot_drivers.srv import GripperDataSet

import numpy as np


class RobotRosNode(Node):
    def __init__(self, robot):
        super().__init__('lerobot_robot_ros_node')
        self.robot = robot
        self.publisher = self.create_publisher(JointState, '/joint_states', 10)
        self.joint_pos_subscriber = self.create_subscription(Jointpos, "/arm/joint_pos", self.set_positions, 10)
        self.enable_service = self.create_service(SetBool, "/arm/enable_motor", self.enable_motor)
        self.gripper_service = self.create_service(GripperDataSet, "/gripper/set_angle", self.gripper)
        self.joint_names = ['joint1', 'joint2', 'joint3', 'joint4', 'joint5', 'joint6']
        self.joint_velocity = [0.0] * len(self.joint_names)
        timer_period = 0.01  # seconds
        self.timer = self.create_timer(timer_period, self.timer_callback)

    def timer_callback(self):
        # 创建 JointState 消息
        msg = JointState()
        msg.header.stamp = self.get_clock().now().to_msg()
        msg.name = self.joint_names
        msg.position = self.robot.get_positions(degree=False).tolist()      # 使用弧度
        msg.velocity = self.joint_velocity
        
        # 发布消息
        # print(msg)
        self.publisher.publish(msg)

    def set_positions(self, pos):
        joints = np.array(pos.joint)
        gripper_pos = self.robot.get_positions()[-1]
        if gripper_pos > np.deg2rad(45):        # 大于45°，张开夹爪
            gripper_pos = 100
        else:
            gripper_pos = 0
        # TODO: 根据夹爪的当前角度，计算百分比

        joints = np.append(joints, gripper_pos)
        print("Set position: ", joints)
        self.robot.set_positions(joints, degree=False)      # 使用的是弧度
    
    def enable_motor(self, req, res):
        # 电机使能与去使能，使能之后电机锁定，去使能之后电机可以拖动
        # ros2 service call /arm/enable_motor std_srvs/srv/SetBool "data: 0"
        if req.data > 0:
            self.robot.enable_motor()
        else:
            self.robot.disable_motor()
        res.success = True
        return res
    
    def gripper(self, req, res):
        # 控制夹爪开合，百分比：0~100，0：合上，100：完全打卡
        # ros2 service call /gripper/set_angle robot_drivers/srv/GripperDataSet "data: 50"
        ang = req.data
        present_pos = self.robot.get_positions()
        present_pos[-1] = ang
        print("Set gripper: ", ang)
        self.robot.set_positions(present_pos, degree=False)
        return res


if __name__ == "__main__":
    from robot_interface import ManipulatorRobot
    from lerobot.common.robot_devices.robots.configs import So101RobotConfig
    rclpy.init()
    config = So101RobotConfig()
    robot = ManipulatorRobot(config)
    robot.connect()

    joint_states_node = RobotRosNode(robot)

    rclpy.spin(joint_states_node)
    joint_states_node.destroy_node()
    rclpy.shutdown()